package com.jboltai.capability.model.claude;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.capability.message.AIMessageRole;
import com.jboltai.capability.message.MessageFinishReason;
import com.jboltai.capability.model.AIEventProcessor;
import com.jboltai.config.AIParamKey;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.EventAttrKey;
import com.jboltai.event.EventErrorType;
import com.jboltai.event.EventState;
import com.jboltai.event.ThinkStatus;
import com.jboltai.event.ai.AIEvent;
import com.jboltai.event.ai.aichat.AIChatEvent;
import com.jboltai.resource.ai.AIResource;
import com.jboltai.util.http.HttpUtil;
import com.jboltai.util.http.RequestInstance;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Objects;

public class ClaudeAIImpl implements AIEventProcessor {
    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);

    public static final ClaudeAIImpl INSTANCE = new ClaudeAIImpl();

    private static final String STOP_FLAG = "{\"type\":\"message_stop\"}";

    public static final String API_URL = "https://api.anthropic.com/v1/messages";

    @Override
    public void process(AIChatEvent event, AIResource resource) {
        //
        RequestInstance request = buildRequest(event, resource);

        event.setAttr(EventAttrKey.REQUEST_INSTANCE, request);

        LOGGER.debug("开始 调用Claude大模型 处理AIChat事件：\n\n\nEvent参数: {}, \n\n\nRequest参数: {}", JSON.toJSONString(event), JSON.toJSONString(request));
        boolean stream = event.getIsStream();
        if (stream) {
            //流式请求
            processEventByStream(event, resource, request);


        } else {
            processEventByNormal(event, resource, request);
        }
    }

    /**
     * 通过正常请求 处理event
     *
     * @param event
     * @param request
     */
    private static void processEventByNormal(AIChatEvent event, AIResource resource, RequestInstance request) {
        request.addJsonParam(AIParamKey.STREAM, false);
        request.asyncSend(response -> {
            try {
                String content = response.body().string();
                LOGGER.debug("AIChatEvent[{}]请求结束:{}", event.getId(), content);
                if (response.code() != 200) {
                    //有异常情况
                    handleFail(event, resource, response.code(), content, null);
                } else {
                    AIMessage message = handleChatSuccess(event, content, false);
                    event.setResult(message);
                    event.setState(EventState.SUCCESS);

                }
            } catch (Exception e) {
                LOGGER.error("AIChatEvent[{}]处理响应结果发生异常:{}", event.getId(), e);
                event.fail(EventErrorType.PROGRAM_EXCEPTION, e);
            } finally {
                event.complete();
            }
        }, e -> {
            LOGGER.error("AIChatEvent[{}]请求失败:{}", event.getId(), e);
            event.fail(EventErrorType.NET_ERROR, e);
        });
    }

    /**
     * 以stream方式处理事件
     *
     * @param event
     * @param resource
     * @param request
     */
    private static void processEventByStream(AIChatEvent event, AIResource resource, RequestInstance request) {
        request.addJsonParam(AIParamKey.STREAM, true);
        request.asyncSendWithSSE(body -> {

            try {
                LOGGER.debug("AIChatEvent[{}]请求成功:{}", event.getId(), body);
                /*
                {"type":"message_start",
                "message":{"id":"msg_bdrk_01SsoygrCimp1kAMBLSbqYxs",
                    "model":"claude-3-haiku-20240307",
                    "type":"message",
                    "role":"assistant",
                    "content":[],
                    "usage":{"input_tokens":10,"output_tokens":2}}}
                 */
                handleChatSuccess(event, body, true);
            } catch (Exception e) {
                event.fail(EventErrorType.PROGRAM_EXCEPTION, e);
            } finally {
                if (event.getState() != EventState.PROCESSING) {
                    event.complete();
                }
            }


        }, (res, e) -> {
            try {
                String response = null;
                if (res != null && res.body() != null) {
                    response = res.body().string();
                }
                handleFail(event, resource, res != null ? res.code() : 0, response, e);
            } catch (Exception ex) {
                ex.printStackTrace();
            }
        });
    }

    /**
     * stream请求时，返回的内容放在"delta" 字段中，非stream请求，返回的内容放在message 字段中
     *
     * @param event
     * @param responseContent
     */
    private static AIMessage handleChatSuccess(AIChatEvent event, String responseContent, boolean stream) {
        /*
        {"id":"msg_bdrk_01XoqaHAXf7xQFoaMqw6yJ5b","model":"claude-3-haiku-20240307","type":"message","role":"assistant",
        "content":[{"type":"text","text":"很高兴认识你!我是一个由人工智能模型开发的对话助手,我可以帮助你解答各种问题,并与你进行友善的交谈。你有什么需要帮助的吗?我会尽我所能为你提供协助。"}],
        "stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":84}}
         */
        /*
        {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"高兴认"}}
         */
        AIMessage message = new AIMessage();
        JSONObject data = JSON.parseObject(responseContent);
        String type = data.getString("type");
        String content = null;
        String finishReason = null;
        switch (type) {
            case "error":
                event.fail(EventErrorType.SERVICE_SERVER_ERROR, data.getJSONObject("error").getString("message"));
                return null;
            case "message_stop":
                event.setState(EventState.SUCCESS);
                return null;
            case "ping":
                return null;
            case "message":
                //非流式返回的消息
                content = data.getJSONArray("content").getJSONObject(0).getString("text");
                finishReason = data.getString("stop_reason");
                event.setPromptTokenCount(data.getJSONObject("usage").getInteger("input_tokens"));
                event.setCompletionTokenCount(data.getJSONObject("usage").getInteger("output_tokens"));
                event.setTotalTokenCount(event.getPromptTokenCount() + event.getCompletionTokenCount());
                break;
            case "message_start":
                event.setPromptTokenCount(data.getJSONObject("message").getJSONObject("usage").getInteger("input_tokens"));
                return null;
            case "content_block_start":
                //流式接口返回的第一帧
                return null;
            case "content_block_stop":
                //最后一个块
                return null;
            case "message_delta":
                //{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":93}}
                event.setCompletionTokenCount(data.getJSONObject("usage").getIntValue("output_tokens"));
                event.setTotalTokenCount(event.getPromptTokenCount() + event.getCompletionTokenCount());
                finishReason = data.getJSONObject("delta").getString("stop_reason");
                break;
            case "content_block_delta":
                //流式数据
                JSONObject delta = data.getJSONObject("delta");
                if (Objects.equals(delta.getString("type"), "thinking_delta")) {
                    message.setThink(delta.getString("thinking"));
                } else {
                    content = delta.getString("text");
                }
                break;
        }
        if (data.containsKey("message")) {
            //流式接口返回的数据
            data = data.getJSONObject("message");
        }

        message.setId(data.getString("id"));
        event.setResultId(message.getId());
        //获取回复
        message.setContent(content);
        event.appendResponse(message.getContent());
//        if (choice.containsKey(contentField)) {
//            message.setContent(choice.getJSONObject(contentField).getString("content"));
//            event.appendResponse(message.getContent());
//        } else {
//            message.setContent("");
//            event.appendResponse("");
//        }
        message.setRole(AIMessageRole.assistant);
        event.setResultRole(message.getRole());
        if (finishReason != null) {
            switch (finishReason) {
                case "end_turn":
                    event.setFinishReason(MessageFinishReason.STOP);
                    break;
                case "max_tokens":
                case "length":
                    event.setFinishReason(MessageFinishReason.LENGTH_LIMIT);
                    break;
                default:
                    event.setFinishReason(MessageFinishReason.OTHER);
            }
        }
        if (StrUtil.isNotBlank(message.getThink()) ) {
            processThinking(event, stream, message.getThink());
        } else if (message.getThink() == null && event.getThinkStatus() == ThinkStatus.RUNNING) {
            event.thinkOver("");
        }
        if (message.getContent() != null) {
            event.success(message);
        }
        return message;
    }
    /**
     * 处理深度思考内容
     * @param event
     * @param stream
     * @param think
     */
    private static void processThinking(AIChatEvent event, boolean stream, String think) {
        if (stream == false) {
            //非流式
            event.thinkOver(think);
            return;
        }
        //流式
        event.thinking(think);
    }
    /**
     * 失败时的处理
     *
     * @param event
     * @param resource
     * @param code
     * @param responseContent
     * @param exception
     */
    private static void handleFail(AIEvent event, AIResource resource, int code, String responseContent, Throwable exception) {
        JSONObject error = null;
        ;
        LOGGER.error("AIEvent[{}]请求失败:{}_{},异常：{}", event.getId(), code, responseContent, exception);
        if (exception == null) {
            //ai服务异常
            switch (code) {
                case 400: {
                    error = JSON.parseObject(responseContent);
                    event.fail(EventErrorType.REQUEST_PARAM_ERROR, error.getJSONObject("error").getString("message"));
                    break;
                }

                case 401:
                    event.fail(EventErrorType.NOT_SUPPORT_MODEL, responseContent);
                    break;
                case 403:
                    event.fail(EventErrorType.API_KEY_INVALID, responseContent);
                    break;
                case 413:
                    event.fail(EventErrorType.AI_EXCEED_TOKEN_LIMIT, responseContent);
                    break;
                case 429: {
                    event.fail(EventErrorType.EXCEED_REQUEST_LIMIT, responseContent);
                    break;
                }
                case 500:
                case 529:
                    event.fail(EventErrorType.SERVICE_SERVER_ERROR);
                    break;
                default:
                    //未知错误
                    event.fail(EventErrorType.UNKNOWN, responseContent);
                    break;
            }
        } else {
            //程序异常
            event.fail(EventErrorType.NET_ERROR, exception);
        }
    }


    private RequestInstance buildRequest(AIChatEvent event, AIResource resource) {
        //资源指定了apiurl就用资源的，否则用默认的
        String apiUrl = StrUtil.isNotBlank(resource.getApiUrl()) ? resource.getApiUrl() : API_URL;
        RequestInstance request = HttpUtil.post(resource.getClient(), apiUrl, event);
        request.addHeader("x-api-key", resource.getApiKey());
        request.addHeader("anthropic-version", "2023-06-01");
        request.setContentTypeToJson();

        request.addJsonParam(AIParamKey.MODEL, event.getModelName());


        if (event.isThinkingEnable()) {
            int thinkingBudget = event.getThinkingBudget() < 1024 ? 1024 : event.getThinkingBudget();
            request.addJsonParam("thinking", new JSONObject().fluentPut("type", "enabled").fluentPut("budget_tokens", thinkingBudget));
            if (event.getMaxTokens() != null) {
                //anthropic的max_tokens是包含thinking的
                request.addJsonParam(AIParamKey.MAX_TOKENS, event.getMaxTokens() + thinkingBudget);
            }
            // claude 深度思考模式不允许传入top k top p
        } else {
            request.addJsonParam("thinking", new JSONObject().fluentPut("type", "disabled"));
            if (event.getMaxTokens() != null) {
                request.addJsonParam(AIParamKey.MAX_TOKENS, event.getMaxTokens());
            }
            if (event.getTopK() != null) { // claude 深度思考模式不允许传入top k
                request.addJsonParam(AIParamKey.TOP_K, event.getTopK());
            }
            if (event.getTopP() != null) {
                request.addJsonParam(AIParamKey.TOP_P, event.getTopP());
            }
            if (event.getTemperature() != null) { //claude思考模式，强制温度为1
                request.addJsonParam(AIParamKey.TEMPERATURE, event.getTemperature() );
            }
        }

        if (event.getMessages() != null) {
            JSONArray messages = new JSONArray(event.getMessages().size());
            JSONObject item = null;
            List<AIMessage> list = event.getMessages();
            for (AIMessage message : list) {
                item = new JSONObject();
                switch (message.getRole()) {
                    case user:
                        item.put("role", "user");
                        break;
                    case assistant:
                        item.put("role", "assistant");
                        break;
                    case system:
                        request.addJsonParam("system", message.getContent());
                        continue;
                    default:
                        continue;
                }
                if (message.getImgUrl() == null) {
                    item.put("content", message.getContent());
                } else {
                    //带图片
//                    JSONArray content = new JSONArray();
//                    content.add(new JSONObject().fluentPut("type", "text").fluentPut("text", message.getContent()));
//                    for (String img : message.getImgUrl()) {
//                        try {
//
//                            content.add(new JSONObject().fluentPut("type", "image_url")
//                                    .fluentPut("image_url", new JSONObject().fluentPut("url", ImgUtil.parseToBase64(img))));
//                        } catch (Exception e) {
//                            LOGGER.error("图片加载失败：{} ", img);
//                        }
//
////                        content.add(new JSONObject().fluentPut("type", "image_url")
////                                .fluentPut("image_url", new JSONObject().fluentPut("url", img)));
//                    }
//                    item.put("content", content);
                }
                messages.add(item);
            }
            request.addJsonParam(AIParamKey.MESSAGES, messages);
        }
        return request;
    }


    @Override
    public void handleCancel(AIChatEvent event, AIResource resource) {
        LOGGER.warn("{}事件取消执行", event.getId());
        RequestInstance requestInstance = event.getAttr(EventAttrKey.REQUEST_INSTANCE);
        requestInstance.cancal();
    }
}
